// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <turbo/log/logging.h>
#include <gtest/gtest.h>

#include <pollux/functions/udf.h>
#include <pollux/testing/functions/prestosql/function_base_test.h>
#include <pollux/type/type.h>
#include <pollux/vector/complex_vector.h>
#include <pollux/vector/selectivity_vector.h>

namespace kumo::pollux {

using namespace kumo::pollux::test;

class SimpleFunctionCallNullFreeTest
    : public functions::test::FunctionBaseTest {};

// Test that function with default contains nulls behavior won't get called when
// inputs are all null.
template <typename T>
struct DefaultContainsNullsBehaviorFunction {
  POLLUX_DEFINE_FUNCTION_TYPES(T);

  bool callNullFree(bool& out, const null_free_arg_type<Array<int32_t>>&) {
    out = true;

    return true;
  }
};

TEST_F(SimpleFunctionCallNullFreeTest, defaultContainsNullsBehavior) {
  register_function<DefaultContainsNullsBehaviorFunction, bool, Array<int32_t>>(
      {"default_contains_nulls_behavior"});

  // Make a vector with some null arrays, and some arrays containing nulls.
  auto array_vector = make_nullable_array_vector<int32_t>(
      {std::nullopt, {{1, std::nullopt}}, {{std::nullopt, 2}}, {{1, 2, 3}}});

  // Check that default contains nulls behavior functions don't get called on a
  // null input.
  auto result = evaluate<SimpleVector<bool>>(
      "default_contains_nulls_behavior(c0)", make_row_vector({array_vector}));
  auto expected = make_nullable_flat_vector<bool>(
      {std::nullopt, std::nullopt, std::nullopt, true});
  assertEqualVectors(expected, result);
}

constexpr int32_t kCallNullable = 0;
constexpr int32_t kCallNullFree = 1;

// Test that function with non-default contains nulls behavior.
// Returns the original array prepended with 0 if callNullable
// is invoked or 1 if callNullFree is invoked.
template <typename T>
struct NonDefaultBehaviorFunction {
  POLLUX_DEFINE_FUNCTION_TYPES(T);

  bool callNullable(
      out_type<Array<int32_t>>& out,
      const arg_type<Array<int32_t>>* input) {
    out.push_back(kCallNullable);

    if (input) {
      for (auto i : *input) {
        if (i.has_value()) {
          out.push_back(*i);
        }
      }
    }

    return true;
  }

  bool callNullFree(
      out_type<Array<int32_t>>& out,
      const null_free_arg_type<Array<int32_t>>& input) {
    out.push_back(kCallNullFree);

    for (auto i : input) {
      out.push_back(i);
    }

    return true;
  }
};

TEST_F(SimpleFunctionCallNullFreeTest, nonDefaultBehavior) {
  register_function<NonDefaultBehaviorFunction, Array<int32_t>, Array<int32_t>>(
      {"non_default_behavior"});

  // Make a vector with a NULL.
  auto arrayVectorWithNull = make_nullable_array_vector<int32_t>(
      {std::nullopt, {{1, 2}}, {{3, 2}}, {{2, 2, 3}}});
  // Make a vector with an array that contains NULL.
  auto arrayVectorContainsNull = make_nullable_array_vector<int32_t>(
      {{4, std::nullopt}, {1, 2}, {3, 2}, {2, 2, 3}});
  // Make a vector that's NULL-free.
  auto arrayVectorNullFree =
      make_array_vector<int32_t>({{4, 5}, {1, 2}, {3, 2}, {2, 2, 3}});

  auto resultWithNull = evaluate<ArrayVector>(
      "non_default_behavior(c0)", make_row_vector({arrayVectorWithNull}));
  auto expectedWithNull = make_array_vector<int32_t>(
      {{kCallNullable},
       {kCallNullable, 1, 2},
       {kCallNullable, 3, 2},
       {kCallNullable, 2, 2, 3}});
  assertEqualVectors(expectedWithNull, resultWithNull);

  auto resultContainsNull = evaluate<ArrayVector>(
      "non_default_behavior(c0)", make_row_vector({arrayVectorContainsNull}));
  auto expectedContainsNull = make_array_vector<int32_t>(
      {{kCallNullable, 4},
       {kCallNullable, 1, 2},
       {kCallNullable, 3, 2},
       {kCallNullable, 2, 2, 3}});
  assertEqualVectors(expectedContainsNull, resultContainsNull);

  auto resultNullFree = evaluate<ArrayVector>(
      "non_default_behavior(c0)", make_row_vector({arrayVectorNullFree}));
  auto expectedNullFree = make_array_vector<int32_t>(
      {{kCallNullFree, 4, 5},
       {kCallNullFree, 1, 2},
       {kCallNullFree, 3, 2},
       {kCallNullFree, 2, 2, 3}});
  assertEqualVectors(expectedNullFree, resultNullFree);
}

template <typename T>
struct DeeplyNestedInputTypeFunction {
  POLLUX_DEFINE_FUNCTION_TYPES(T);

  bool callNullFree(
      bool& out,
      const null_free_arg_type<
          Variadic<Row<Array<int32_t>, Map<int32_t, int32_t>>>>&) {
    out = true;

    return true;
  }
};

TEST_F(SimpleFunctionCallNullFreeTest, deeplyNestedInputType) {
  register_function<
      DeeplyNestedInputTypeFunction,
      bool,
      Variadic<Row<Array<int32_t>, Map<int32_t, int32_t>>>>(
      {"deeply_nested_input_type"});

  vector_size_t size = 11;
  auto arrayVector1 = make_array_vector<int32_t>(
      size,
      [&](vector_size_t /* row */) { return 1; },
      [&](vector_size_t idx) { return idx; },
      // First array is NULL.
      [&](vector_size_t row) { return row == 0; },
      // Second array contains NULL.
      [&](vector_size_t idx) { return idx == 0; });
  auto mapVector1 = make_map_vector<int32_t, int32_t>(
      size,
      [&](vector_size_t /* row */) { return 1; },
      [&](vector_size_t idx) { return idx; },
      [&](vector_size_t idx) { return idx; },
      // Third map is NULL.
      [&](vector_size_t row) { return row == 2; },
      // Fourth map has a NULL value.
      [&](vector_size_t idx) { return idx == 2; });
  auto rowVector1 = make_row_vector(
      {"array", "map"},
      {arrayVector1, mapVector1},
      // Fifth row is NULL.
      [&](vector_size_t row) { return row == 4; });
  auto arrayVector2 = make_array_vector<int32_t>(
      size,
      [&](vector_size_t /* row */) { return 1; },
      [&](vector_size_t idx) { return idx; },
      // Sixth array is NULL.
      [&](vector_size_t row) { return row == 5; },
      // Seventh array contains NULL.
      [&](vector_size_t idx) { return idx == 5; });
  auto mapVector2 = make_map_vector<int32_t, int32_t>(
      size,
      [&](vector_size_t /* row */) { return 1; },
      [&](vector_size_t idx) { return idx; },
      [&](vector_size_t idx) { return idx; },
      // Eigth map is NULL.
      [&](vector_size_t row) { return row == 7; },
      // Ninth map has a NULL value.
      [&](vector_size_t idx) { return idx == 7; });
  auto rowVector2 = make_row_vector(
      {"array", "map"},
      {arrayVector2, mapVector2},
      // Tenth row is NULL
      [&](vector_size_t row) { return row == 9; });

  auto result = evaluate<SimpleVector<bool>>(
      "deeply_nested_input_type(c0, c1)",
      make_row_vector({rowVector1, rowVector2}));
  auto expected = make_nullable_flat_vector<bool>(
      {std::nullopt,
       std::nullopt,
       std::nullopt,
       std::nullopt,
       std::nullopt,
       std::nullopt,
       std::nullopt,
       std::nullopt,
       std::nullopt,
       std::nullopt,
       true});
  assertEqualVectors(expected, result);
}

} // namespace kumo::pollux
